fix: jax reducers returning incorrect output values or lengths #3464

ikrommyd · 2025-04-14T00:14:37Z

Needs more work and I'd appreciate any help @ianna @pfackeldey.
I'm adding a test that tests all the reducers via parametrization.

Definitely needs #3457

This PR should fix #3456, #3462, #3463 and #3465

ikrommyd · 2025-04-14T00:26:27Z

I'm currently skipping the argmin argmax tests because of #3463 but that should change.

ikrommyd · 2025-04-14T02:28:37Z

awkward/src/awkward/_connect/jax/reducers.py

Lines 296 to 299 in c93da2f

    
           # See issue https://github.com/google/jax/issues/9296 
        
           result = jax.numpy.exp( 
        
               jax.ops.segment_sum(jax.numpy.log(array.data), parents.data) 
        
           )

So we can't take the product of an array with negative numbers? The logarithm will just NaN the output

ikrommyd · 2025-04-14T03:13:59Z

awkward/src/awkward/_connect/jax/reducers.py

Lines 296 to 299 in c93da2f

# See issue https://github.com/google/jax/issues/9296

result = jax.numpy.exp(

jax.ops.segment_sum(jax.numpy.log(array.data), parents.data)

)

So we can't take the product of an array with negative numbers? The logarithm will just NaN the output

I may have fixed that in 9874810. I still haven't made ak.any to behave like the cpu case when segments have zeros though.

ianna · 2025-04-14T09:29:15Z

@ikrommyd - impressive! There are only 5 tests that fails: 3 for ak.any, 1 for ak.count, and 1 for ak.sum. It looks like the latter is failing with a boolean dtype. Perhaps, return should make an int?

ianna

Thanks for looking into it!

ianna · 2025-04-14T09:43:06Z

src/awkward/_connect/jax/reducers.py

@@ -261,7 +328,7 @@ def apply(
        if array.dtype.kind == "M":
            raise TypeError(f"cannot compute the sum (ak.sum) of {array.dtype!r}")



@ikrommyd - if we want to allow sum of boolean, I think, we should view the data as integers here:

if array.dtype == np.bool_: data = array.data.astype(jax.numpy.int32) else: data = array.data

ikrommyd · 2025-04-14T14:36:28Z

So I made ci pass but there are a few things that are to be done for sure

The code needs refactoring, I don't like how it looks at all, it's very hacky
The reducer tests should definitely test more input array cases
We need to make sure I'm not breaking something that should work but is untested due to the smaller amount of tests for the jax backend

ianna

@ikrommyd - just some minor comments. I agree, JAX backend needs to be thoroughly tested. Perhaps, our fellow could take over? @pfackeldey - when do we discuss his project? Thanks.

ianna · 2025-04-14T09:51:43Z

src/awkward/_connect/jax/reducers.py

@@ -68,7 +131,7 @@ def segment_argmin(data, segment_ids):
 class ArgMin(JAXReducer):
    name: Final = "argmin"
    needs_position: Final = True
-    preferred_dtype: Final = np.int64
+    preferred_dtype: Final = np.float64


argmin returns the index (i.e. position) of the minimum value. Indices are always integers, not floating-point numbers.

ianna · 2025-04-14T09:56:24Z

src/awkward/_connect/jax/reducers.py

@@ -125,7 +191,7 @@ def segment_argmax(data, segment_ids):
 class ArgMax(JAXReducer):
    name: Final = "argmax"
    needs_position: Final = True
-    preferred_dtype: Final = np.int64
+    preferred_dtype: Final = np.float64


the same argument here - indices are always integers, not floating-point numbers.

… of differentiability

ikrommyd added 2 commits April 13, 2025 19:09

add tests and use outlength

c5d85e3

rename test and skip if jax is not installed

1f74feb

ikrommyd mentioned this pull request Apr 14, 2025

ak.any yields True for empty rows in float arrays with the jax backend #3465

Open

ikrommyd temporarily deployed to docs April 14, 2025 00:28 — with GitHub Actions Inactive

ikrommyd added 2 commits April 13, 2025 21:02

count and argmax should be returning integers

4215ef9

maybe make ak.any cpu and jax behavior match?

7e681af

ikrommyd temporarily deployed to docs April 14, 2025 02:27 — with GitHub Actions Inactive

skip product checks with negative values

51a0945

ikrommyd temporarily deployed to docs April 14, 2025 02:46 — with GitHub Actions Inactive

segment product that deals with negative and zero values too

9874810

ikrommyd temporarily deployed to docs April 14, 2025 03:23 — with GitHub Actions Inactive

simplify the function

28432d6

ikrommyd temporarily deployed to docs April 14, 2025 03:39 — with GitHub Actions Inactive

my first attempt at apply_positional_corrections for jax

6423929

ikrommyd temporarily deployed to docs April 14, 2025 07:35 — with GitHub Actions Inactive

Merge branch 'main' into ikrommyd/jax-backend-reducers

74a788f

ikrommyd temporarily deployed to docs April 14, 2025 09:04 — with GitHub Actions Inactive

ianna reviewed Apr 14, 2025

View reviewed changes

.at is better than concatenate

bfae901

ikrommyd temporarily deployed to docs April 14, 2025 14:00 — with GitHub Actions Inactive

fix cases like ak.any([[0, 0], [1, 0], [2, 3, 4]]) maybe?

32cf4e9

ikrommyd temporarily deployed to docs April 14, 2025 14:20 — with GitHub Actions Inactive

fix sum and count with bools?

c35c7d4

ikrommyd temporarily deployed to docs April 14, 2025 14:33 — with GitHub Actions Inactive

ikrommyd mentioned this pull request Apr 14, 2025

feat: enable jax backend for virtual arrays #3451

Draft

Merge branch 'main' into ikrommyd/jax-backend-reducers

1892261

ianna reviewed Apr 14, 2025

View reviewed changes

ianna temporarily deployed to docs April 14, 2025 18:01 — with GitHub Actions Inactive

Merge branch 'main' into ikrommyd/jax-backend-reducers

8178f4e

ikrommyd temporarily deployed to docs April 15, 2025 19:11 — with GitHub Actions Inactive

Merge branch 'main' into ikrommyd/jax-backend-reducers

5bc2976

ikrommyd temporarily deployed to docs April 16, 2025 18:09 — with GitHub Actions Inactive

ikrommyd added 2 commits April 20, 2025 20:44

make jax reducers return the same dtype as the numpy ones at the cost…

bdfa953

… of differentiability

Merge branch 'main' into ikrommyd/jax-backend-reducers

8338000

ikrommyd deployed to docs April 21, 2025 02:01 — with GitHub Actions View deployment

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

fix: jax reducers returning incorrect output values or lengths #3464

fix: jax reducers returning incorrect output values or lengths #3464

ikrommyd commented Apr 14, 2025 •

edited

Loading

ikrommyd commented Apr 14, 2025

ikrommyd commented Apr 14, 2025

ikrommyd commented Apr 14, 2025

ianna commented Apr 14, 2025

ianna left a comment

ianna Apr 14, 2025

ikrommyd commented Apr 14, 2025 •

edited

Loading

ianna left a comment

ianna Apr 14, 2025

ianna Apr 14, 2025

		@@ -261,7 +328,7 @@ def apply(
		if array.dtype.kind == "M":
		raise TypeError(f"cannot compute the sum (ak.sum) of {array.dtype!r}")

fix: jax reducers returning incorrect output values or lengths #3464

Are you sure you want to change the base?

fix: jax reducers returning incorrect output values or lengths #3464

Conversation

ikrommyd commented Apr 14, 2025 • edited Loading

ikrommyd commented Apr 14, 2025

ikrommyd commented Apr 14, 2025

ikrommyd commented Apr 14, 2025

ianna commented Apr 14, 2025

ianna left a comment

Choose a reason for hiding this comment

ianna Apr 14, 2025

Choose a reason for hiding this comment

ikrommyd commented Apr 14, 2025 • edited Loading

ianna left a comment

Choose a reason for hiding this comment

ianna Apr 14, 2025

Choose a reason for hiding this comment

ianna Apr 14, 2025

Choose a reason for hiding this comment

ikrommyd commented Apr 14, 2025 •

edited

Loading

ikrommyd commented Apr 14, 2025 •

edited

Loading